import os
import pickle
import bpy
import mathutils
import numpy as np

# === 路径设置 ===
PKL_DIR = './mocap_output_full/mocap'  # 修正路径
SMPLX_MODEL_PATH = './extra_data/smpl/SMPLX_NEUTRAL.pkl'
EXPORT_PATH = './output_fbx/smplx_animation.fbx'

# === 创建输出目录 ===
os.makedirs(os.path.dirname(EXPORT_PATH), exist_ok=True)

# === 先安装 smplx 到 Blender 的 Python 环境 ===
try:
    import smplx
    import torch
    print("Dependencies already installed")
except ImportError:
    print("Installing dependencies...")
    import subprocess
    import sys
    
    # 获取 Blender 的 Python 可执行文件路径
    python_exe = sys.executable
    
    # 安装必要的包
    subprocess.check_call([python_exe, "-m", "pip", "install", "torch", "smplx"])
    
    # 重新导入
    import smplx
    import torch
    print("Dependencies installed successfully")

print("Loading SMPLX model...")

# === 加载 SMPLX 模型 ===
MODEL = smplx.create(
    SMPLX_MODEL_PATH,
    model_type='smplx', 
    gender='neutral', 
    use_face_contour=True,
    ext='pkl'  # 指定使用 pkl 格式
)

print("SMPLX model loaded successfully")

# === 清理场景 ===
bpy.ops.object.select_all(action='SELECT')
bpy.ops.object.delete(use_global=False)

print("Scene cleared")

# === 获取动画文件 ===
pkl_files = sorted([f for f in os.listdir(PKL_DIR) if f.endswith('.pkl')])
print(f"Found {len(pkl_files)} animation frames")

if len(pkl_files) == 0:
    print("Warning: No .pkl files found in mocap_output_full/mocap directory")
    exit()

# === 设置动画范围 ===
bpy.context.scene.frame_start = 1
bpy.context.scene.frame_end = len(pkl_files)
bpy.context.scene.frame_set(1)

print("Reading first frame to understand data structure...")

# === 读取第一帧数据来了解结构 ===
with open(os.path.join(PKL_DIR, pkl_files[0]), 'rb') as f:
    first_frame_data = pickle.load(f)

print(f"Data keys: {first_frame_data.keys()}")

# 根据 FrankMocap 的输出格式解析数据
demo_type = first_frame_data.get('demo_type', 'body')
smpl_type = first_frame_data.get('smpl_type', 'smplx')
pred_output_list = first_frame_data.get('pred_output_list', [])

if len(pred_output_list) == 0:
    print("No prediction outputs found in the first frame")
    exit()

# 使用第一个人的预测结果
pred_output = pred_output_list[0]
if pred_output is None:
    print("First person prediction is None")
    exit()

print(f"Demo type: {demo_type}, SMPL type: {smpl_type}")
print(f"Prediction keys: {pred_output.keys()}")

# === 创建初始 SMPLX Mesh ===
def extract_smpl_params(pred_output):
    """从预测结果中提取 SMPL 参数"""
    
    # 根据数据格式提取参数
    if 'pred_body_pose' in pred_output:
        body_pose = pred_output['pred_body_pose']
        if body_pose.shape == (24, 3, 3):  # rotation matrices
            # 转换为角轴表示
            from scipy.spatial.transform import Rotation as R
            body_pose = R.from_matrix(body_pose.reshape(-1, 3, 3)).as_rotvec().reshape(-1)
            body_pose = body_pose[3:]  # 排除全局朝向
        elif body_pose.shape == (23, 3):  # 已经是角轴格式
            body_pose = body_pose.reshape(-1)
        else:
            body_pose = body_pose.reshape(-1)
    else:
        body_pose = np.zeros(69)  # 23 joints * 3
    
    if 'pred_global_orient' in pred_output:
        global_orient = pred_output['pred_global_orient']
        if global_orient.shape == (1, 3, 3):
            from scipy.spatial.transform import Rotation as R
            global_orient = R.from_matrix(global_orient.reshape(3, 3)).as_rotvec()
        else:
            global_orient = global_orient.reshape(3)
    else:
        global_orient = np.zeros(3)
    
    betas = pred_output.get('pred_betas', np.zeros(10))
    if betas.shape[0] > 10:
        betas = betas[:10]
    
    # 手部和面部参数
    left_hand_pose = pred_output.get('pred_left_hand_pose', np.zeros(45))
    right_hand_pose = pred_output.get('pred_right_hand_pose', np.zeros(45))
    jaw_pose = pred_output.get('pred_jaw_pose', np.zeros(3))
    expression = pred_output.get('pred_expression', np.zeros(10))
    
    return {
        'body_pose': torch.tensor(body_pose, dtype=torch.float32).unsqueeze(0),
        'global_orient': torch.tensor(global_orient, dtype=torch.float32).unsqueeze(0),
        'betas': torch.tensor(betas, dtype=torch.float32).unsqueeze(0),
        'left_hand_pose': torch.tensor(left_hand_pose.reshape(-1), dtype=torch.float32).unsqueeze(0),
        'right_hand_pose': torch.tensor(right_hand_pose.reshape(-1), dtype=torch.float32).unsqueeze(0),
        'jaw_pose': torch.tensor(jaw_pose.reshape(-1), dtype=torch.float32).unsqueeze(0),
        'expression': torch.tensor(expression.reshape(-1), dtype=torch.float32).unsqueeze(0)
    }

# 检查是否有预计算的顶点
if 'pred_vertices_smpl' in pred_output:
    print("Using pre-computed vertices from pkl file")
    use_precomputed_vertices = True
else:
    print("Computing vertices using SMPL-X model")
    use_precomputed_vertices = False

# 创建基础网格
if use_precomputed_vertices:
    verts = pred_output['pred_vertices_smpl']
    if verts.shape[0] == 1:
        verts = verts.squeeze(0)
    faces = MODEL.faces
else:
    params = extract_smpl_params(pred_output)
    output = MODEL(**params, return_verts=True)
    verts = output.vertices.detach().cpu().numpy().squeeze()
    faces = MODEL.faces

print(f"Generated mesh with {len(verts)} vertices and {len(faces)} faces")

# === 创建 Mesh 对象 ===
def create_mesh(name, verts, faces):
    mesh = bpy.data.meshes.new(name)
    mesh.from_pydata(verts.tolist(), [], faces.tolist())
    mesh.update()
    obj = bpy.data.objects.new(name, mesh)
    bpy.context.collection.objects.link(obj)
    return obj

smplx_obj = create_mesh("SMPLX", verts, faces)
smplx_obj.select_set(True)
bpy.context.view_layer.objects.active = smplx_obj

print("Base mesh created")

# === 创建形状键动画 ===
# 添加基础形状键
basis_key = smplx_obj.shape_key_add(name="Basis")
basis_key.interpolation = 'KEY_LINEAR'

# 为每一帧创建形状键
for frame_idx, fname in enumerate(pkl_files):
    print(f"Processing frame {frame_idx + 1}/{len(pkl_files)}: {fname}")
    
    with open(os.path.join(PKL_DIR, fname), 'rb') as f:
        data = pickle.load(f)

    pred_output_list = data.get('pred_output_list', [])
    if len(pred_output_list) == 0 or pred_output_list[0] is None:
        print(f"Skipping frame {frame_idx + 1}: no valid prediction")
        continue
        
    pred_output = pred_output_list[0]  # 使用第一个人的数据
    
    # 获取当前帧的顶点
    if use_precomputed_vertices and 'pred_vertices_smpl' in pred_output:
        verts = pred_output['pred_vertices_smpl']
        if verts.shape[0] == 1:
            verts = verts.squeeze(0)
    else:
        try:
            params = extract_smpl_params(pred_output)
            output = MODEL(**params, return_verts=True)
            verts = output.vertices.detach().cpu().numpy().squeeze()
        except Exception as e:
            print(f"Error computing vertices for frame {frame_idx + 1}: {e}")
            continue
    
    # 创建当前帧的形状键
    key_name = f"Frame_{frame_idx:04d}"
    shape_key = smplx_obj.shape_key_add(name=key_name)
    
    # 设置形状键的顶点位置
    for i, v in enumerate(verts):
        shape_key.data[i].co = mathutils.Vector(v)
    
    # 设置动画关键帧
    current_frame = frame_idx + 1
    bpy.context.scene.frame_set(current_frame)
    
    # 清除所有形状键的值
    for key in smplx_obj.data.shape_keys.key_blocks:
        if key.name != "Basis":
            key.value = 0.0
            key.keyframe_insert(data_path="value", frame=current_frame)
    
    # 只激活当前帧的形状键
    shape_key.value = 1.0
    shape_key.keyframe_insert(data_path="value", frame=current_frame)
    
    # 为前一帧和后一帧设置为0（创建清晰的过渡）
    if current_frame > 1:
        bpy.context.scene.frame_set(current_frame - 1)
        shape_key.value = 0.0
        shape_key.keyframe_insert(data_path="value", frame=current_frame - 1)
    
    if current_frame < len(pkl_files):
        bpy.context.scene.frame_set(current_frame + 1)
        shape_key.value = 0.0
        shape_key.keyframe_insert(data_path="value", frame=current_frame + 1)

print("Animation processing completed")

# === 设置线性插值 ===
if smplx_obj.data.shape_keys:
    for key_block in smplx_obj.data.shape_keys.key_blocks:
        key_block.interpolation = 'KEY_LINEAR'

# === 确保对象有动画数据 ===
bpy.context.scene.frame_set(1)
smplx_obj.location = (0, 0, 0)
smplx_obj.keyframe_insert(data_path="location", frame=1)
smplx_obj.keyframe_insert(data_path="location", frame=len(pkl_files))

print("Setting up animation curves...")

# === 导出 FBX 时包含动画 ===
print("Exporting FBX with animation...")

# 确保选中对象
bpy.ops.object.select_all(action='DESELECT')
smplx_obj.select_set(True)
bpy.context.view_layer.objects.active = smplx_obj

# 导出 FBX，包含动画数据
bpy.ops.export_scene.fbx(
    filepath=EXPORT_PATH,
    use_selection=True,
    apply_unit_scale=True,
    bake_space_transform=True,
    object_types={'MESH'},
    use_mesh_modifiers=True,
    mesh_smooth_type='FACE',
    use_tspace=True,
    use_custom_props=False,
    add_leaf_bones=False,
    primary_bone_axis='Y',
    secondary_bone_axis='X',
    use_armature_deform_only=False,
    armature_nodetype='NULL',
    bake_anim=True,
    bake_anim_use_all_bones=False,
    bake_anim_use_nla_strips=False,
    bake_anim_use_all_actions=False,
    bake_anim_force_startend_keying=True,
    bake_anim_step=1.0,
    bake_anim_simplify_factor=1.0,
    path_mode='AUTO',
    embed_textures=False,
    batch_mode='OFF'
)

print(f"Exported FBX to {EXPORT_PATH}")
print("Export completed successfully!")

# === 保存 Blender 文件用于调试 ===
blend_path = EXPORT_PATH.replace('.fbx', '.blend')
bpy.ops.wm.save_as_mainfile(filepath=blend_path)
print(f"Also saved Blender file to {blend_path} for debugging")
